package gorpc

import (
	"errors"
	"fmt"
	"sync"
	"sync/atomic"
	"time"

	pb "gorpc/pb"

	"github.com/gorilla/websocket"
	"google.golang.org/protobuf/proto"
)

// RPCSession 表示一个 RPC 会话
type RPCSession struct {
	mu sync.RWMutex

	// 连接相关
	conn      *websocket.Conn
	closed    atomic.Bool
	closeOnce sync.Once
	done      chan struct{}

	// 请求相关
	nextSeqID atomic.Int32
	handlers  sync.Map // map[int32]*RPCCallContext

	// 配置相关
	readTimeout  time.Duration
	writeTimeout time.Duration
	maxMsgSize   int64

	// 回调函数
	onConnected    func()
	onDisconnected func(err error)
	onRequest      func(inMsg *RPCMessage, outMsg *RPCMessage) error
	onNotify       func(msg *RPCMessage)
}

// NewRPCSession 创建一个新的 RPC 会话
func NewRPCSession(conn *websocket.Conn) *RPCSession {
	s := &RPCSession{
		conn:         conn,
		readTimeout:  60 * time.Second,
		writeTimeout: 10 * time.Second,
		maxMsgSize:   32 * 1024 * 1024, // 32MB
		done:         make(chan struct{}),
	}
	return s
}

// SetReadTimeout 设置读取超时时间
func (s *RPCSession) SetReadTimeout(timeout time.Duration) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.readTimeout = timeout
}

// SetWriteTimeout 设置写入超时时间
func (s *RPCSession) SetWriteTimeout(timeout time.Duration) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.writeTimeout = timeout
}

// SetMaxMessageSize 设置最大消息大小
func (s *RPCSession) SetMaxMessageSize(size int64) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.maxMsgSize = size
}

// SetConnectedCallback 设置连接建立回调
func (s *RPCSession) SetConnectedCallback(cb func()) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.onConnected = cb
}

// SetDisconnectedCallback 设置连接断开回调
func (s *RPCSession) SetDisconnectedCallback(cb func(error)) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.onDisconnected = cb
}

// SetRequestCallback 设置请求处理回调
func (s *RPCSession) SetRequestCallback(cb func(inMsg *RPCMessage, outMsg *RPCMessage) error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.onRequest = cb
}

// SetNotifyCallback 设置通知回调
func (s *RPCSession) SetNotifyCallback(cb func(msg *RPCMessage)) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.onNotify = cb
}

// Start 启动会话
func (s *RPCSession) Start() error {
	if s.closed.Load() {
		return errors.New("session is closed")
	}

	if s.onConnected != nil {
		s.onConnected()
	}

	go s.readLoop()
	return nil
}

// Done 返回一个通道，当会话结束时关闭
func (s *RPCSession) Done() <-chan struct{} {
	return s.done
}

// Close 关闭会话
func (s *RPCSession) Close() error {
	var err error
	s.closeOnce.Do(func() {
		s.closed.Store(true)

		// 关闭所有等待中的请求
		s.handlers.Range(func(key, value interface{}) bool {
			if ctx, ok := value.(*RPCCallContext); ok {
				ctx.Cancel()
			}
			s.handlers.Delete(key)
			return true
		})

		// 关闭连接
		if s.conn != nil {
			err = s.conn.Close()
		}

		// 触发断开回调
		if s.onDisconnected != nil {
			s.onDisconnected(err)
		}

		// 关闭 done 通道
		close(s.done)
	})
	return err
}

// IsClosed 检查会话是否已关闭
func (s *RPCSession) IsClosed() bool {
	return s.closed.Load()
}

// CallMethod 使用消息对象执行RPC调用
func (s *RPCSession) CallMethod(inMsg *RPCMessage, outMsg *RPCMessage, timeout time.Duration) error {
	if s.closed.Load() {
		return errors.New("session is closed")
	}

	if inMsg == nil || outMsg == nil {
		return errors.New("invalid message")
	}

	// 创建请求上下文
	seqID := s.nextSeqID.Add(1)
	callCtx := NewRPCCallContext(seqID, timeout)
	s.handlers.Store(seqID, callCtx)
	defer s.handlers.Delete(seqID)

	// 序列化输入消息
	coder := NewMessageProtobufCoder()
	msgData, err := coder.Encode(inMsg)
	if err != nil {
		return fmt.Errorf("encode request message error: %w", err)
	}

	// 发送请求
	packet := &pb.Packet{
		PacketType: pb.PacketType_Request,
		Method:     inMsg.GetMessageID(),
		Context:    msgData,
		Request: &pb.RequestPacket{
			SeqId: seqID,
		},
	}

	if err := s.writePacket(packet); err != nil {
		return err
	}

	// 等待响应
	response, err := callCtx.Wait()
	if err != nil {
		return err
	}

	// 解码响应消息
	if len(response) > 0 {
		respMsg, err := coder.Decode(response)
		if err != nil {
			return fmt.Errorf("decode response message error: %w", err)
		}
		// 复制响应消息的所有属性到输出消息
		outMsg.SetMessageID(respMsg.GetMessageID())
		outMsg.SetMessage(respMsg.GetMessage())
		outMsg.SetProps(respMsg.GetProps())
		outMsg.SetLevel(respMsg.GetLevel())
	}
	return nil
}

// Notify 发送通知消息
func (s *RPCSession) Notify(msg *RPCMessage) error {
	if s.closed.Load() {
		return errors.New("session is closed")
	}

	if msg == nil {
		return errors.New("message is nil")
	}

	// 序列化消息
	coder := NewMessageProtobufCoder()
	msgData, err := coder.Encode(msg)
	if err != nil {
		return fmt.Errorf("encode notify message error: %w", err)
	}

	packet := &pb.Packet{
		PacketType: pb.PacketType_RpcIssue,
		Method:     msg.GetMessageID(),
		Context:    msgData,
	}

	return s.writePacket(packet)
}

// readLoop 读取消息循环
func (s *RPCSession) readLoop() {
	defer s.Close()

	for {
		if s.closed.Load() {
			return
		}

		// 设置读取超时
		if s.readTimeout > 0 {
			s.conn.SetReadDeadline(time.Now().Add(s.readTimeout))
		}

		// 读取消息
		_, message, err := s.conn.ReadMessage()
		if err != nil {
			if !s.closed.Load() {
				s.Close()
			}
			return
		}

		// 解码消息
		packet := &pb.Packet{}
		if err := proto.Unmarshal(message, packet); err != nil {
			continue
		}

		// 处理消息
		go s.handlePacket(packet)
	}
}

// handlePacket 处理收到的消息包
func (s *RPCSession) handlePacket(packet *pb.Packet) {
	switch packet.PacketType {
	case pb.PacketType_Request:
		s.handleRequest(packet)
	case pb.PacketType_Response:
		s.handleResponse(packet)
	case pb.PacketType_RpcIssue:
		s.handleNotify(packet)
	}
}

// handleRequest 处理请求消息
func (s *RPCSession) handleRequest(packet *pb.Packet) {
	if s.onRequest == nil {
		return
	}

	// 解码输入消息
	coder := NewMessageProtobufCoder()
	inMsg, err := coder.Decode(packet.Context)
	if err != nil {
		fmt.Printf("Failed to decode request message: %v\n", err)
		return
	}
	inMsg.SetMessageID(packet.Method)

	// 创建输出消息
	outMsg := NewMessage()

	// 处理请求
	err = s.onRequest(inMsg, outMsg)

	// 编码输出消息
	var respData []byte
	if err == nil {
		respData, err = coder.Encode(outMsg)
		if err != nil {
			fmt.Printf("Failed to encode response message: %v\n", err)
			err = fmt.Errorf("encode response error: %w", err)
		}
	}

	// 发送响应
	respPacket := &pb.Packet{
		PacketType: pb.PacketType_Response,
		Method:     packet.Method,
		Response: &pb.ResponsePacket{
			SeqId: packet.Request.SeqId,
		},
	}

	if err != nil {
		errMsg := NewMessage()
		errMsg.SetErrorMessage(err.Error())
		if data, encodeErr := coder.Encode(errMsg); encodeErr == nil {
			respPacket.Context = data
		}
	} else {
		respPacket.Context = respData
	}

	s.writePacket(respPacket)
}

// handleResponse 处理响应消息
func (s *RPCSession) handleResponse(packet *pb.Packet) {
	if packet.Response == nil {
		return
	}

	seqID := packet.Response.SeqId
	value, ok := s.handlers.Load(seqID)
	if !ok {
		return
	}

	callCtx, ok := value.(*RPCCallContext)
	if !ok {
		return
	}

	// 解码响应消息
	var err error
	if len(packet.Context) > 0 {
		coder := NewMessageProtobufCoder()
		msg, decodeErr := coder.Decode(packet.Context)
		if decodeErr == nil {
			if errStr := msg.GetProp("error"); errStr != "" {
				err = errors.New(errStr)
				callCtx.Complete(nil, err)
			} else {
				callCtx.Complete(packet.Context, nil)
			}
		} else {
			err = fmt.Errorf("decode response error: %w", decodeErr)
			callCtx.Complete(nil, err)
		}
	} else {
		callCtx.Complete(nil, nil)
	}
}

// handleNotify 处理通知消息
func (s *RPCSession) handleNotify(packet *pb.Packet) {
	if s.onNotify == nil {
		return
	}

	coder := NewMessageProtobufCoder()
	msg, err := coder.Decode(packet.Context)
	if err != nil {
		fmt.Printf("Failed to decode notify message: %v\n", err)
		return
	}
	msg.SetMessageID(packet.Method)
	s.onNotify(msg)
}

// writePacket 发送消息包
func (s *RPCSession) writePacket(packet *pb.Packet) error {
	if s.closed.Load() {
		return errors.New("session is closed")
	}

	// 设置写入超时
	if s.writeTimeout > 0 {
		s.conn.SetWriteDeadline(time.Now().Add(s.writeTimeout))
	}

	// 编码消息
	data, err := proto.Marshal(packet)
	if err != nil {
		return err
	}

	// 发送消息
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.conn.WriteMessage(websocket.BinaryMessage, data)
}

// GetID 获取会话ID
func (s *RPCSession) GetID() string {
	return s.conn.RemoteAddr().String()
}

// GetRemoteAddr 获取远程地址
func (s *RPCSession) GetRemoteAddr() string {
	return s.conn.RemoteAddr().String()
}

// GetLocalAddr 获取本地地址
func (s *RPCSession) GetLocalAddr() string {
	return s.conn.LocalAddr().String()
}

// IsConnected 检查是否已连接
func (s *RPCSession) IsConnected() bool {
	return !s.closed.Load()
}
